from torch.nn import Linear, Module

class ActionModel(Module):
    def __init__(self, statnum, musclenum):
        
        self.f1=Linear(statnum, 30)
        self.fh=[Linear(30,30) for i in range(6)]
        self.f2=Linear(30,musclenum*2)

    def forward(self,x):
        x=self.f1(x).relu()
        for i in self.fh:
            x+=i(x).relu()
        x=self.f2(x)
        x=Ten.connect([x.cut(i*2,i*2+2).softmax() for i in range(len(x)//2)])
        return x

    def choice(self,x):
        v = self.forward(x).data
        a=[]
        # print(v)
        for i in range(len(v)):
            if i%2==1:
                continue
            v2=v[i:i+2]
            a.append(v2.index(random.choices(v2,v2)[0]))
        Operator.clean()
        return a

    def optimize(self,k=0.01):
        self.f1.grad_descent_zero(k)
        for i in self.fh:
            i.grad_descent_zero(k)
        self.f2.grad_descent_zero(k)